package com.suxin.ratelimit.v2.handler;

import lombok.AllArgsConstructor;
import lombok.Data;
import lombok.NoArgsConstructor;
import lombok.extern.slf4j.Slf4j;

import java.time.LocalDateTime;
import java.time.ZoneId;
import java.util.Date;
import java.util.Objects;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.locks.ReentrantLock;

/**
 * @author Tang
 * @classname StandRateLimitHandler
 * @description [  ]
 * @date 2022/2/10 10:04
 */
@Slf4j
public class StandRateLimitHandler implements RateLimitHandler {

    private final ConcurrentHashMap<String, RateLimit> resultMap = new ConcurrentHashMap<>();

    private final ReentrantLock lock = new ReentrantLock();

    @Override
    public boolean doHandler(Rate rate, String key) {
        lock.lock();
        try {
            RateLimit rateLimit = resultMap.get(key);
            Date date = new Date();
            if (Objects.isNull(rateLimit) || rateLimit.getEndDate().compareTo(date) < 0) {
                Date endDate = Date.from(LocalDateTime.ofInstant(date.toInstant(), ZoneId.systemDefault())
                        .plus(rate.getDuration(), rate.getUnit()).atZone(ZoneId.systemDefault()).toInstant());
                rateLimit = new RateLimit(1, endDate);
                resultMap.put(key, rateLimit);
                return false;
            }
            // 结束时间 > 当前时间
            if (rateLimit.getEndDate().compareTo(date) > 0 && rateLimit.getCount() < rate.getLimit()) {
                rateLimit.setCount(rateLimit.getCount() + 1);
            } else if (rateLimit.getCount() >= rate.getLimit()) {
                return true;
            }
        } catch (Exception e) {
            e.printStackTrace();
        } finally {
            lock.unlock();
        }
        return false;
    }

    @Data
    @AllArgsConstructor
    @NoArgsConstructor
    static class RateLimit {
        /**
         * 访问次数
         */
        int count;

        /**
         * 结束时间
         */
        Date endDate;
    }

}